#include "spellcheck.h"
#include "BaselineSpellChecker.h"
#include "PerWordSpellChecker.h"
#include "DefaultSpellChecker.h"
#include "trie.h"
#include "exception.h"
#include "string_utils.h"
#include "BigramCollection.h"

#include <string>
#include <memory>
#include <iostream>
#include <iomanip>
#include <fstream>
#include <vector>
#include <stdexcept>
#include <algorithm>
#include <cstdlib>
#include <cstring>

using std::string;
using std::istream;
using std::ostream;
using std::vector;

const int PRECISION = 8;

void processRequests(const NSpellCheck::ISpellChecker& checker, istream& in, ostream& out) {
    using namespace NSpellCheck;
    string input;
    out << std::setiosflags(std::ios::fixed) << std::setprecision(PRECISION);
    while (std::getline(in, input)) {
        Phrase phrase;
        parsePhrase(input, &phrase);
        vector<Variation> suggestions;
        checker.process(phrase, &suggestions);
        out << phrase;
        for (vector<Variation>::const_iterator i = suggestions.begin(); i != suggestions.end(); ++i) {
            out << "\t" << *i;
        }
        out << std::endl;
    }
}

void bulkProcessRequests(const NSpellCheck::ISpellChecker& checker, istream& in, ostream& out) {
    using namespace NSpellCheck;
    vector<string> inputLines;
    string input;
    while (std::getline(in, input)) {
        inputLines.push_back(input);
    }
    const int numInputs = inputLines.size();
    vector<vector<Variation> > answers(numInputs);
    std::cerr << "Preprocessing done, running main job..." << std::endl;
#pragma omp parallel for
    for (int i = 0; i < numInputs; ++i) {
        Phrase phrase;
        parsePhrase(inputLines[i], &phrase);
        checker.process(phrase, &answers[i]);
    }
    std::cerr << "Finalizing..." << std::endl;
    out << std::setiosflags(std::ios::fixed) << std::setprecision(PRECISION);
    for (int i = 0; i < numInputs; ++i) {
        out << inputLines[i];
        for (vector<Variation>::const_iterator it = answers[i].begin(); it != answers[i].end(); ++it) {
            out << "\t" << *it;
        }
        out << "\n";
    }
}

void readFrequencies(istream& in, Trie<int>* trie) {
    string input;
    while (std::getline(in, input)) {
        vector<string> tokens;
        split(input, '\t', &tokens);
        ASSERT(tokens.size() == 2);
        trim(tokens[0]);
        trim(tokens[1]);
        int thisFreq = atoi(tokens[1].c_str());
        int* prevFreq = trie->addWord(tokens[0]).first;
        *prevFreq += thisFreq;
    }
}

void readBigrams(istream& in, NSpellCheck::BigramCollection* bigrams) {
    string input;
    while (std::getline(in, input)) {
        vector<string> tokens;
        split(input, '\t', &tokens);
        ASSERT(tokens.size() == 2);
        trim(tokens[0]);
        trim(tokens[1]);
        vector<string> words;
        split(tokens[0], ' ', &words);
        ASSERT(words.size() == 2);
        int freq = atoi(tokens[1].c_str());
        if (freq < 30) {
            break;
        }
        bigrams->addBigram(words[0], words[1], freq);
    }
}

void analyzeSentences(istream& in, Trie<int>* unigrams, NSpellCheck::BigramCollection* bigrams) {
    using namespace NSpellCheck;
    string input;
    while (std::getline(in, input)) {
        killPunctuation(input);
        vector<string> tokens;
        split(input, ' ', &tokens);
        vector<const int*> iters;
        iters.push_back(BigramCollection::SENTENCE_START);
        for (vector<string>::iterator i = tokens.begin(); i != tokens.end(); ++i) {
            int* freq = unigrams->addWord(*i).first;
            ++ *freq;
            iters.push_back(freq);
        }
        iters.push_back(BigramCollection::SENTENCE_END);
        for (size_t i = 0; i + 1 < iters.size(); ++i) {
            bigrams->addBigram(iters[i], iters[i+1]);
        }
    }
}

NSpellCheck::ISpellChecker* createSpellChecker(int argc, char* argv[]) {
    using namespace NSpellCheck;
    if (argc == 1) {
        return new BaselineSpellChecker();
    }
    string spellCheckerName = argv[1];
    if (spellCheckerName == "baseline") {
        return new BaselineSpellChecker();
    } else if (spellCheckerName == "perword") {
        std::auto_ptr<Trie<int> > unigrams(new Trie<int>);
        std::fstream unigramsFile(argv[2]);
        readFrequencies(unigramsFile, unigrams.get());
        return new PerWordSpellChecker(unigrams.release());
    } else if (spellCheckerName == "def") {
        std::auto_ptr<Trie<int> > unigrams(new Trie<int>);
        std::fstream unigramsFile(argv[2]);
        readFrequencies(unigramsFile, unigrams.get());
        std::auto_ptr<BigramCollection> bigrams(new BigramCollection(*unigrams));
        std::fstream bigramsFile(argv[3]);
        readBigrams(bigramsFile, bigrams.get());
        return new DefaultSpellChecker(unigrams.release(), bigrams.release());
    } else if (spellCheckerName == "def2") {
        std::auto_ptr<Trie<int> > unigrams(new Trie<int>);
        std::auto_ptr<BigramCollection> bigrams(new BigramCollection(*unigrams));
        std::fstream text(argv[2]);
        analyzeSentences(text, unigrams.get(), bigrams.get());
        return new DefaultSpellChecker(unigrams.release(), bigrams.release());
    } else {
        throw std::invalid_argument(spellCheckerName + ": unknown spell checker name");
    }
}

int main(int argc, char* argv[]) {
    bool interactive = false;
    if (argc > 1 && !strcmp(argv[1], "-int")) {
        interactive = true;
        --argc;
        ++argv;
        std::cerr << "Loading, please wait... ";
    }
    std::auto_ptr<NSpellCheck::ISpellChecker> spellChecker(createSpellChecker(argc, argv));
    if (interactive) {
        std::cerr << "Ready!" << std::endl;
        processRequests(*spellChecker, std::cin, std::cout);
    } else {
        bulkProcessRequests(*spellChecker, std::cin, std::cout);
    }
    return 0;
}
